import cv2
import numpy as np
from image_extractor import ImageVectorExtractor
import config
from matplotlib import pyplot as plt
import tools
import numba
from multiprocessing import Pool, Queue

INPUT_BLOCK_SIZE = 10


def move_block(x, y, block, new_cloud, block_speeds):
    # 取平均速度
    speedx_matrix = block_speeds[:, :, 0]
    speed_x = speedx_matrix.mean()

    speedy_matrix = block_speeds[:, :, 1]
    speed_y = speedy_matrix.mean()

    nx = int(x + speed_x)
    ny = int(y + speed_y)

    # nx = int(x + 100)
    # ny = int(y)

    if (nx > 0 and nx + INPUT_BLOCK_SIZE < new_cloud.shape[1]) and (
                    ny > 0 and ny + INPUT_BLOCK_SIZE < new_cloud.shape[0]):
        newBlock = new_cloud[ny:ny + INPUT_BLOCK_SIZE, nx:nx + INPUT_BLOCK_SIZE]
        radarValue = block

        if newBlock.shape != radarValue.shape:
            return

        for i in range(newBlock.shape[0]):
            for k in range(newBlock.shape[1]):
                if newBlock[i][k] == 0:
                    newBlock[i][k] = radarValue[i][k]
                else:
                    value = max(newBlock[i][k], radarValue[i][k])
                    newBlock[i][k] = value

                    # if newBlock.shape == block[:, :, -1:].shape:
                    #     img[ny:ny + INPUT_BLOCK_SIZE, nx:nx + INPUT_BLOCK_SIZE] = block[:, :, -1:]


def pre(full_speed_matrix, cloud):
    new_cloud = np.zeros((full_speed_matrix.shape[0], full_speed_matrix.shape[1]), np.uint32)

    x = 0
    y = 0
    count = 0

    while True:
        block_values = cloud[y:y + INPUT_BLOCK_SIZE, x:x + INPUT_BLOCK_SIZE]
        block_speeds = full_speed_matrix[y:y + INPUT_BLOCK_SIZE, x:x + INPUT_BLOCK_SIZE]

        if block_values.mean() > 0:
            # print("x:%d,y:%d,max:%d,min:%d, mean:%f" % (
            #     x, y, block_values.max(), block_values.min(), block_values.mean()))

            move_block(x, y, block_values, new_cloud, block_speeds)

            count = count + 1

        x = x + 5
        if x >= cloud.shape[1]:
            x = 0
            y += 5
            if y >= cloud.shape[0]:
                break

    # cv2.imshow('new', new_cloud)
    # cv2.waitKey()

    return new_cloud


def speed_vectors():
    speed_matrix_list = []

    shape = ()
    for i in range(30, 48):
        imageVectorExactor = ImageVectorExtractor(config.RADAR_IMAGE_PATH + '%02d.png' % (i,),
                                                  config.RADAR_IMAGE_PATH + '%02d.png' % (i + 1,))

        # imageVectorExactor = ImageVectorExtractor(config.BLOCK + '%02d.png' % (i,),
        #                                           config.BLOCK + '%02d.png' % (i + 1,))

        print("process :", imageVectorExactor.path1 + "," + imageVectorExactor.path2)

        imageVectorExactor.match()
        # imageVectorExactor.showImg1MatchedPoints()
        imageVectorExactor.extract_vectors()
        speed_matrix_list.append(imageVectorExactor.matchedPointVectors)
        # speed_matrix_list.append(imageVectorExactor.speed_matrix)

        shape = imageVectorExactor.img1.shape

    speed_matrix_list.reverse()

    matrix = np.zeros((shape[0], shape[1], 2), np.float32)

    for matchedPointVectors in speed_matrix_list:
        for vector in matchedPointVectors:
            x = vector[0]
            y = vector[1]
            matrix[y, x] = vector[2]

    np.save('matrix.npy', matrix)

    # data = matrix[:, :, 0:1]
    # data = data.reshape(shape)
    #
    # plt.imshow(data)
    # plt.colorbar()
    # plt.show()


def findNearerMatchedPoint(point, matchedPointVectors, num=3):
    global count
    x = np.array([e[0] for e in matchedPointVectors])
    y = np.array([e[1] for e in matchedPointVectors])
    distance = (x - point[0]) ** 2 + (y - point[1]) ** 2

    idx = np.argsort(distance)

    nearer = [matchedPointVectors[i][2] for i in idx[:3]]
    # if count % 1000 == 0:
    #     print(count)
    # count += 1

    return nearer


def nearer_vector(nearer):
    '''邻近点的速度计算平均速度'''

    speed_x_total = 0
    speed_y_total = 0
    for p in nearer:
        # 从临近点速度计算出平均速度
        speed_x = p[0]
        speed_y = p[1]

        speed_x_total += speed_x
        speed_y_total += speed_y

    speed_x_avg = speed_x_total / len(nearer)
    speed_y_avg = speed_y_total / len(nearer)

    return speed_x_avg, speed_y_avg


count = 0

grid_size = 5


def fill_block_speed(args):
    global processed_block_count
    processed_block_count += 1
    x, y, matchedPointVectors, speed_matrix = args
    # x = point[0]
    # y = point[1]
    block_values = speed_matrix[y:y + grid_size, x:x + grid_size]
    center_x = x
    center_y = y
    nearer = findNearerMatchedPoint((center_x, center_y), matchedPointVectors)
    avg = nearer_vector(nearer)
    block_values[:] = avg

    print("total:%d,cur:%d" % (fill_block_total, processed_block_count))

    return block_values


fill_block_total = 0
processed_block_count = 0


# @numba.jit
def calc_not_match_vectors(speed_matrix, matchedPointVectors):
    '''计算未匹配的点的运功矢量'''

    print('计算未匹配的点的运功矢量')
    # for y in range(speed_matrix.shape[0]):
    #     for x in range(speed_matrix.shape[1]):
    #         vector = speed_matrix[y][x]
    #         if vector[0] == 0:  # 未匹配
    #             print((x, y))
    #             nearer = findNearerMatchedPoint((x, y), matchedPointVectors)
    #             avg = nearer_vector(nearer)
    #             vector[0] = avg[0]
    #             vector[1] = avg[1]

    global fill_block_total
    x = 0
    y = 0
    arg_list = []
    while True:
        arg_list.append((x, y, matchedPointVectors, speed_matrix))
        x = x + grid_size
        if x >= speed_matrix.shape[1]:
            x = 0
            y += grid_size
            if y >= speed_matrix.shape[0]:
                break
    fill_block_total = len(arg_list) / 8
    print("start,total:%d......................" % (fill_block_total,))
    pool = Pool(8)
    results = pool.map(fill_block_speed, arg_list)

    for args, result in zip(arg_list, results):
        x, y, matchedPointVectors, speed_matrix = args
        block_values = speed_matrix[y:y + grid_size, x:x + grid_size]
        block_values[:] = result

    print(".........................end")


import time


def fix_full_speed_vectors():
    print('fix_full_speed_vectors')
    matrix = np.load('matrix.npy')
    x_idx, y_idx = np.where(matrix[:, :, 0] > 0)
    idx = np.concatenate((x_idx.reshape(-1, 1), y_idx.reshape(-1, 1)), axis=1)
    matchedPointVectors = [(x, y, matrix[y, x]) for y, x in idx]
    matchedPointVectors.sort(key=lambda point: point[0])

    start = time.time()
    calc_not_match_vectors(matrix, matchedPointVectors)
    end = time.time()

    print("cost:%f:" % (end - start))
    np.save('full_speed.npy', matrix)

    # data = matrix[:, :, 0:1]
    # data = data.reshape(matrix.shape[0], matrix.shape[1])
    # print(data.shape)
    # plt.imshow(data)
    # plt.colorbar()
    # plt.show()


if __name__ == '__main__':
    # speed_vectors()
    # fix_full_speed_vectors()
    # exit(0)

    full_speed_matrix = np.load('full_speed.npy')

    # data = full_speed_matrix[:, :, 0:1]
    # data = data.reshape(full_speed_matrix.shape[0], full_speed_matrix.shape[1])
    # print(data.shape)
    # plt.imshow(data)
    # plt.colorbar()
    # plt.show()

    # cloud = cv2.imread(config.BLOCK + '60.png', cv2.COLOR_BGR2GRAY)

    cloud = np.load(config.RADAR_MATRIX_PATH + '48.npy')
    # cloud = np.load(config.BLOCK + '48.npy')
    cloud = np.clip(cloud, 0, None)
    # image = np.clip(cloud, 0, None)
    # image = image / 4
    # image = image.astype(np.uint8)
    # cv2.imwrite('out/00.jpg', image)

    i = 0
    while i < 10:
        print("pre img %d" % (i,))
        cloud = pre(full_speed_matrix, cloud)
        np.save('out/new_%02d.npy' % (i,), cloud)

        image = cloud.copy()
        image = image / 4
        image = image.astype(np.uint8)
        cv2.imwrite('out/new_%02d.jpg' % (i,), image)
        i += 1
        # cv2.waitKey()
